# author:巇橙
# data：2023.3.29

import pandas as pd
import math


# 定义信息熵函数
def get_Ent(data_arr):
	class_arr = []  # 存储类别
	num_arr = []  # 存储类别的数目
	new_data_arr = set(data_arr)
	for i in new_data_arr:
		sum = 0
		class_arr.append(i)
		for j in range(0, len(data_arr)):
			if i == data_arr[j]:
				sum += 1
		num_arr.append(sum)
	# 对输入的列表的单列元素种类及数量进行统计得到列表class_arr和列表num_arr，两个列表为key-value映射关系
	# class_arr = [a1,a2,a3]
	# num_arr = [2,3,4]
	num_all = 0  # 初始化个数
	Ent = 0  # 初始化信息熵
	for i in num_arr:
		num_all += i
	for i in num_arr:
		Ent -= (i / num_all) * math.log2(i / num_all)
	return Ent


# 定义根据类别划分数据
def class_split(data_arr, pre_arr):
	# 此处的data_arr和pre_arr的关系依然为映射列表
	Ent = get_Ent(pre_arr)
	new_data_arr = set(data_arr)
	for i in new_data_arr:
		sum = 0  # 初始化统计变量
		class_arr = []  # 初始化类别变量
		for j in range(0, len(data_arr)):
			if i == data_arr[j]:
				sum += 1
				class_arr.append(pre_arr[j])
		new_Ent = get_Ent(class_arr)
		Ent -= (sum / len(data_arr)) * new_Ent
	return Ent


# 定义获取数据函数
def get_data(data):
	data = pd.read_excel(data)
	data_A = data.iloc[:9, 0].tolist()  # a1
	data_B = data.iloc[:9, 1].tolist()  # a2
	data_C = data.iloc[:9, 2].tolist()  # a3
	data_D = data.iloc[:9, 3].tolist()  # 标签
	return data_A, data_B, data_C, data_D


# 定义连续值处理函数
def get_coiled(data_arr, pre_arr):
	sort_arr = sorted(set(data_arr))
	cent_arr = []  # 初始化二分类中心列表
	Ent_arr = []  # 初始化信息增益列表
	for i in range(0, len(sort_arr) - 1):
		cent_arr.append((sort_arr[i] + sort_arr[i + 1]) / 2)
	for i in cent_arr:
		Ent = get_Ent(pre_arr)
		sum = 0
		min_arr = []  # 初始化小于中心值列表
		max_arr = []  # 初始化大于中心值列表
		for j in range(0, len(data_arr)):
			if i < data_arr[j]:
				sum += 1  # 将大于中心值的数据累加
				max_arr.append(pre_arr[j])
			elif i > data_arr[j]:
				min_arr.append(pre_arr[j])
		new_Ent_max = get_Ent(max_arr)
		new_Ent_min = get_Ent(min_arr)
		print(new_Ent_max)
		print(new_Ent_min)
		Ent -= (sum / len(data_arr) * new_Ent_max) + ((len(data_arr) - sum) / len(data_arr)) * new_Ent_min  # 计算信息增益
		Ent_arr.append(Ent)
	return Ent_arr, cent_arr  # 返回信息增益列表及二分类中心列表


# 主函数入口
if __name__ == "__main__":
	data_A1, data_A2, data_A3, data_pre = get_data("data2.xlsx")
	print("信息熵为：", round(get_Ent(data_pre), 3))
	print("a1的信息增益为：", round(class_split(data_A1, data_pre), 3))
	print("a2的信息增益为：", round(class_split(data_A2, data_pre), 3))
	Ent_arr, cent_arr = get_coiled(data_A3, data_pre)
	for i in range(0, len(Ent_arr)):
		print("中心为", cent_arr[i], "的信息增益为：", round(Ent_arr[i], 3))
